"""
词向量
"""
import os

from gensim.models import word2vec, KeyedVectors
import numpy as np

from .settings import WV_MODEL_PATH
from .util import timing


class WordVec:
    def __init__(self, wv_model=None):
        """
        Args
        ----
        wv_model : gensim加载的词向量模型对象
        """
        if wv_model is None:
            self._wv = self.load_wv()
        else:
            self._wv = wv_model
        vec_size = self._wv.vector_size  # 向量的维度
        self.NULL_VEC = np.zeros((1, vec_size))

    def load_wv(self):
        """加载词向量
        """
        wv = KeyedVectors.load_word2vec_format(WV_MODEL_PATH, binary=True)
        return wv

    def doc2vec(self, doc):
        """分词后的文本转为向量
        Args
        ----
        doc : iterable, 迭代单词 

        Returns
        -------
        vec : np.ndarray, 一维向量
        """
        vec = np.array([self._wv[w] for w in doc if w in self._wv])
        # 词语都不在与训练的词向量中，则句子向量的各维度均为0
        if len(vec) == 0:
            vec = self.NULL_VEC
        return vec.mean(axis=0)

    def docs2vec(self, docs):
        """多个分词后的文本转向量矩阵

        Args
        ----
        docs : iterable, 迭代单个doc

        Returns
        -------
        vec : np.ndarray, vec矩阵
        """
        return np.array([self.doc2vec(doc) for doc in docs])

